"""
Discover fusion genes based on RNA sequencing data.

Usage:
  phusion detect <reads>... [--output-dir=PATH]
  phusion blacklist <fusion_report>...
  phusion filter <fusion_report> [--blacklist=PATH] [--min-flank-len=N]
  phusion rank <fusion_report>...
  phusion validate <fusion> <reads>...
  phusion visualize <fusion_report>...

Options:
  -h --help      Show this screen.
"""

from __future__ import print_function
import sys, subprocess, docopt, re, datetime, pwd, os, operator, itertools
from collections import defaultdict
from pypette import mkdir, shell, read_fasta, revcomplement, GenomicFeatures
from pypette import info, error


##################
# PHUSION DETECT #
##################

def detect_fusions(read_paths, out_dir=None, anchor_len=25):
	
	sample_name = os.path.basename(read_paths[0])
	sample_name = re.sub('(_[12])?\.(fastq|fasta|fq|fa).*', '', sample_name,
		flags=re.I)
	
	if not out_dir: out_dir = 'fusions'
	out_dir = re.sub('/+$', '', out_dir) + '/' + sample_name
	mkdir(out_dir)
		
	info('Starting fusion analysis on %s...' % sample_name)
	
	info('Discarding reads that align against the transcriptome...')
	txome_unaligned_path = out_dir + '/txome_unaligned.fq.gz'
	shell('bowtie2 -p8 --score-min L,0,-0.2 '
		'-x ~/tools/bowtie2-indexes/homo_sapiens/ensembl_68 --un-gz %s -U %s '
		'> /dev/null' % (txome_unaligned_path, ' '.join(read_paths)))
	
	info('Discarding reads that align against the genome...')
	genome_unaligned_path = out_dir + '/genome_unaligned.fq.gz'
	shell('bowtie2 -p8 --score-min L,0,-0.2 '
		'-x ~/tools/bowtie2-indexes/homo_sapiens/hg19 --un-gz %s -U %s '
		'> /dev/null' % (genome_unaligned_path, txome_unaligned_path))

	info('Splitting unaligned reads into %d bp anchors...' % anchor_len)
	anchor_5p_path = out_dir + '/anchors_5p.fq.gz'
	anchor_3p_path = out_dir + '/anchors_3p.fq.gz'
	shell('python ~/pypette/split_fasta_reads.py %s %d %s %s' %
		(genome_unaligned_path, anchor_len, anchor_5p_path, anchor_3p_path))
	
	info('Aligning anchors against the genome...')
	anchor_alignments_path = out_dir + '/anchor_alignments.sam'
	shell('bowtie2 -f -p8 --score-min L,0,0 --no-unal '
		'-x ~/tools/bowtie2-indexes/homo_sapiens/hg19 -1 %s -2 %s > %s' %
		(anchor_5p_path, anchor_3p_path, anchor_alignments_path))
	
	discordant = []
	
	prev = ['']
	for line in open(anchor_alignments_path):
		if line[0] in '@#': continue
		tokens = line[:-1].split('\t')
		
		if tokens[0] == prev[0]:
			if tokens[2] == prev[2] and \
				abs(int(tokens[3]) - int(prev[3])) < 100000:
				continue
			
			# Check that the anchors align uniquely
			if int(tokens[4]) == 0 or int(prev[4]) == 0: continue
			
			# Mates are discordant, store them in a list.
			# Each pair is represented as a 7-tuple in the order
			# (chr_1, strand_1, pos_1, chr_2, strand_2, pos_2, sequence).
			# Positions are 1-based and represent read starts.
			full_seq = prev[0][prev[0].find('_')+1:]
			pair = [prev[2], '-' if int(prev[1]) & 16 else '+', int(prev[3]),
				tokens[2], '-' if int(tokens[1]) & 16 else '+', int(tokens[3]),
				full_seq]
			if pair[1] == '-': pair[2] += anchor_len - 1
			if pair[4] == '-': pair[5] += anchor_len - 1
				
			discordant.append(pair)
		else:
			prev = tokens
		
	info('Found %d discordant mate pairs.' % len(discordant))

	# Now we scan through all discordant mate pairs and look for the most likely
	# breakpoint location.
	chromosomes = read_fasta(
		'/data/csb/tools/bowtie2-indexes/homo_sapiens/hg19.fa')
	
	fusions = defaultdict(list)
	
	for k, pair in enumerate(discordant):
		if pair[0] == 'chrM' or pair[3] == 'chrM': continue
		
		seq = pair[6]
		full_len = len(seq)
		
		# Extract the flanking sequences from the chromosome sequences.
		if pair[1] == '+':
			left_grch = chromosomes[pair[0]][pair[2]-1:pair[2]+full_len-1]
		else:
			left_grch = revcomplement(
				chromosomes[pair[0]][pair[2]-full_len:pair[2]])
		
		if pair[4] == '+':
			right_grch = chromosomes[pair[3]][
				pair[5]+anchor_len-full_len-1:pair[5]+anchor_len-1]
		else:
			right_grch = revcomplement(chromosomes[pair[3]][
				pair[5]-anchor_len:pair[5]-anchor_len+full_len])
			
		# Check that the read sequence is not too homologous on either side
		# of the breakpoint.
		left_match = float(sum([seq[i] == left_grch[i]
			for i in range(full_len - anchor_len, full_len)])) / anchor_len
		right_match = float(sum([seq[i] == right_grch[i]
			for i in range(anchor_len)])) / anchor_len

		if left_match >= 0.7 or right_match >= 0.7: continue
		
		# Identify the breakpoint location that minimizes the number of
		# nucleotide mismatches between the read and breakpoint flanks.
		potential_breakpoints = range(anchor_len, full_len - anchor_len + 1)
		mismatches = [0] * len(potential_breakpoints)
		for k, br in enumerate(potential_breakpoints):
			grch_chimera = left_grch[:br] + right_grch[br:]
			mismatches[k] = sum([seq[i] != grch_chimera[i]
				for i in range(full_len)])
			
		# The best breakpoint placement cannot have more than N mismatches.
		least_mismatches = min(mismatches)
		if least_mismatches > 2: continue
		
		# Items in "best_breakpoints" represent the RNA-seq read nucleotide
		# after which the breakpoint occurs. Note that there can be multiple
		# best breakpoints if a microhomology exists at the fusion junction.
		best_breakpoints = [potential_breakpoints[k]
			for k, m in enumerate(mismatches) if m == min(mismatches)]
		
		# Lists "left_breakpoints" and "right_breakpoints" list in genomic
		# coordinates the last nucleotides directly next to the breakpoint.
		# A valid fusion junction is produced by taking elements with the
		# same index from both lists.
		if pair[1] == '+':
			left_breakpoints = [pair[2] + bp - 1 for bp in best_breakpoints]
		else:
			left_breakpoints = [pair[2] - bp + 1 for bp in best_breakpoints]
			
		if pair[4] == '+':
			right_breakpoints = [pair[5] + bp - 1 for bp in best_breakpoints]
		else:
			right_breakpoints = [pair[5] - bp + 1 for bp in best_breakpoints]
			
		left_breakpoints = tuple(left_breakpoints)
		right_breakpoints = tuple(right_breakpoints)
		
		# A putative fusion gene is represented as a 6-tuple that encodes
		# the breakpoint location on both sides of the junction.
		fusion = (pair[0], pair[1], left_breakpoints,
			pair[3], pair[4], right_breakpoints)
		
		# Now that we know the exact fusion breakpoint, we mark mismatches
		# with a lower case nucleotide and augment the read
		# sequence with a | symbol to denote the junction.
		br = abs(pair[2] - left_breakpoints[0]) + 1
		grch_chimera = left_grch[:br] + right_grch[br:]
		seq = ''.join([nuc if grch_chimera[k] == nuc else nuc.lower()
			for k, nuc in enumerate(seq)])
		seq = seq[:br] + '|' + seq[br:]
		pair[6] = seq
		
		# As we cannot assume strand specific sequencing, all fusion genes
		# are reoriented so that the "left" breakpoint always has the lower
		# coordinate or chromosome. Without this reorienting step, one fusion
		# might be represented by two entries in the "fusions" dictionary due
		# to reads originating from both strands.
		if fusion[0] > fusion[3] or \
			(fusion[0] == fusion[3] and fusion[2] > fusion[5]):
			fusion = (
				fusion[3], '-' if fusion[4] == '+' else '+', fusion[5][::-1],
				fusion[0], '-' if fusion[1] == '+' else '+', fusion[2][::-1])

		fusions[fusion].append(pair)
		

	
	# A valid fusion candidate must have at least three unique reads.
	for fusion, reads in fusions.items():
		unique_reads = list(set([tuple(r)[0:6] for r in reads]))
		if len(unique_reads) < 3:
			del fusions[fusion]
			
			
			
	info('Found %d fusion candidates.' % len(fusions))
	
			
	# Determine nearby genomic features for fusion breakpoints.
	genomic_features = GenomicFeatures.from_bed('/data/csb/organisms'
		'/homo_sapiens/ensembl_68/genes.bed')
			
	
	
	sorted_fusions = sorted(fusions.iteritems(), key=lambda x: len(x[1]),
		reverse=True)
	
	report_path = out_dir + '/fusions.tsv'
	report = open(report_path, 'w')

	report.write('CHROM\tSTRAND\tPOSITION\tFEATURES\t \t'
		'CHROM\tSTRAND\tPOSITION\tFEATURES\t \tNUM_READS\tREAD_SEQUENCES\n')
	for fusion, reads in sorted_fusions:
		report.write('%s\t%s\t%d\t%s\t \t%s\t%s\t%d\t%s\t \t%d\t%s\n' % (
			fusion[0], fusion[1], fusion[2][0], ', '.join(['%s (%d)' % f
			for f in genomic_features.near(fusion[0], fusion[2][0], 100000)]),
			fusion[3], fusion[4], fusion[5][0], ', '.join(['%s (%d)' % f
			for f in genomic_features.near(fusion[3], fusion[5][0], 100000)]),
			len(reads), ', '.join([read[6] for read in reads])))
	report.close()

	
			
#####################
# PHUSION BLACKLIST #
#####################

def create_fusion_blacklist(fusion_report_paths):
	
	blacklist_radius = r = 5000 
	
	for path in fusion_report_paths:
		fusion_report = open(path)
		for line in fusion_report:
			if not line.startswith('chr'): continue
			tokens = line[:-1].split('\t')
			print('%s\t%d\t%d\t-\t-\t+' % (tokens[0], int(tokens[2]) - r - 1,
				int(tokens[2]) + r))
			print('%s\t%d\t%d\t-\t-\t+' % (tokens[5], int(tokens[7]) - r - 1,
				int(tokens[7]) + r))
		fusion_report.close()








##################
# PHUSION FILTER #
##################

def filter_fusions(fusion_report_path, region_blacklist_path=None,
	min_flank_len=0):
	
	info('Filtering fusions in sample %s...' % fusion_report_path)
	
	# Blacklist of gene or fusion names that often show up as false positives.
	gene_blacklist = [
		'IGL.+-IGL.+', 'IGK.+-IGK.+', 'IGLV\d', 'IGLC\d', 'IGLL\d',
		'IGLJ\d', 'IGLVV', '^ACT.+-ACT.+', 'RPS.+-RPS.+', 'GAPDH.+-GAPDH.+',
		'RN5-8S', 'RN18S', 'RN28S', 'HLA-', 'LINC00273', 'RP11-']
	gene_blacklist = [re.compile(pattern) for pattern in gene_blacklist]

	# Blacklist of genomic regions that are often involved in gene fusions.
	region_blacklist = []
	if region_blacklist_path:
		region_blacklist = GenomicFeatures.from_bed(region_blacklist_path)
	
	for line in open(fusion_report_path):
		if not line.startswith('chr'):
			print(line)
			continue
		
		tokens = line[:-1].split('\t')
		
		# Check that neither of the breakpoints is within a blacklisted region.
		if region_blacklist.near(tokens[0], int(tokens[2]), 0): continue
		if region_blacklist.near(tokens[5], int(tokens[7]), 0): continue
			
		# Check that reads cover a sufficient amount of nucleotides on both
		# sides of the breakpoint.
		read_seqs = tokens[11].split(', ')
		longest_left_flank = max([seq.find('|') for seq in read_seqs])
		longest_right_flank = max(
			[len(seq) - seq.find('|') - 1 for seq in read_seqs])
		if min([longest_left_flank, longest_right_flank]) < min_flank_len:
			continue
		
		# Check that the fusion does not involve any blacklisted genes.
		# Take out the silly (ENSG839239) stuff. FIXME: Remove this.
		tokens[3] = re.sub(' \(ENSG\d+\)', '', tokens[3])
		tokens[8] = re.sub(' \(ENSG\d+\)', '', tokens[8])
		
		nearby_features_1 = re.finditer(r'(\S+) \((\d+)\)', tokens[3])
		nearby_features_2 = re.finditer(r'(\S+) \((\d+)\)', tokens[8])
		
		nearby_features_1 = [m.group(1) for m in nearby_features_1
			if abs(int(m.group(2))) <= 0]
		nearby_features_2 = [m.group(1) for m in nearby_features_2
			if abs(int(m.group(2))) <= 0]
		
		# Remove any fusion candidates that only involve one gene.
		if len(set(nearby_features_1) & set(nearby_features_2)) > 0:
			continue
		
		# Check if the fusion involves blacklisted genes.
		skip = False
		for pair in itertools.product(nearby_features_1, nearby_features_2):
			if any([rx.search('%s-%s' % pair)!=None for rx in gene_blacklist]):
				skip = True
				break
		if skip: continue

		# All clear with this fusion.
		print(line)
			
		
		
		
	






################
# PHUSION RANK #
################

def rank_fusions(fusion_report_paths, out_prefix=None):
	
	if not out_prefix: out_prefix = '/home/annalam/fusions'
	#mkdir(out_dir)
	
	sample_names = [re.sub('\.tsv$', '', path, flags=re.I)
		for path in fusion_report_paths]
	S = len(sample_names)
	
	gene_pair_rearrangements = defaultdict(lambda: [[] for s in range(S)])
	
	for s, report_path in enumerate(fusion_report_paths):
		report = open(report_path)
		for line in report:
			if not line.startswith('chr'): continue
			
			tokens = line[:-1].split('\t')
			
			nearby_features_1 = re.finditer(r'(\S+) \((\d+)\)', tokens[3])
			nearby_features_2 = re.finditer(r'(\S+) \((\d+)\)', tokens[8])
			
			nearby_features_1 = [m.group(1) for m in nearby_features_1
				if abs(int(m.group(2))) <= 0]
			nearby_features_2 = [m.group(1) for m in nearby_features_2
				if abs(int(m.group(2))) <= 0]

			# Place the supporting reads in the sample-fusion matrix.
			read_seqs = tokens[11].split(', ')
			for pair in itertools.product(nearby_features_1, nearby_features_2):
					gene_pair_rearrangements[pair][s] += read_seqs

		report.close()
		
	# Write out the spreadsheet and read sequenecs.
	fusion_matrix = open(out_prefix + '_matrix.tsv', 'w')
	fusion_matrix.write('FUSION\tNUM_POSITIVE\t%s\n' % '\t'.join(sample_names))
	
	fusion_reads = open(out_prefix + '_reads.tsv', 'w')
	
	format_str = '\t%s' * S + '\n'
	for pair, reads in sorted(gene_pair_rearrangements.iteritems(),
		key=lambda x: sum([len(r) > 0 for r in x[1]]), reverse=True):
		fusion_matrix.write('%s-%s\t%d' % (pair[0], pair[1],
			sum([len(r) > 0 for r in reads])))
		fusion_matrix.write(format_str % tuple([len(r) for r in reads]))
		
		fusion_reads.write('%s-%s:\n' % (pair[0], pair[1]))
		for s, sample_reads in enumerate(reads):
			for read in sample_reads:
				fusion_reads.write('%d: %s\n' % (s+1, read))
		fusion_reads.write('\n')
				
	fusion_matrix.close()
	fusion_reads.close()
	

		
		
	




####################
# PHUSION VALIDATE #
####################

def validate_fusion(fusion, read_paths, out_dir=None):
	
	if not out_dir: out_dir = '/home/annalam/fusion_validations/%s' % fusion
	mkdir(out_dir)
	
	gene_1, gene_2 = fusion.split('-')
	
	chromosomes = read_fasta(
		'/data/csb/tools/bowtie2-indexes/homo_sapiens/hg19.fa')

	exons_1 = []
	exons_2 = []
	
	gtf_file = open('/data/csb/organisms/homo_sapiens/ensembl_68/'
		'Homo_sapiens.GRCh37.68.gtf')
	for line in gtf_file:
		tokens = line.split('\t')
		if not tokens[2] == 'exon': continue
		if tokens[1] in ['nonsense_mediated_decay', 'retained_intron']: continue
		if tokens[0].startswith('HSCHR'): continue
		if tokens[0].startswith('GL'): continue
		if tokens[0].startswith('HG'): continue

		gene_name = re.search('gene_name "(.+?)"', line).group(1)
		if gene_name != gene_1 and gene_name != gene_2: continue
		
		tx_id = re.search('transcript_id "(.+?)"', line).group(1)
		exon_num = re.search('gene_name "(.+?)"', line).group(1)
		
		
		if gene_name == gene_1:
			exons_1.append(
				('chr' + tokens[0], tokens[6], int(tokens[3]), int(tokens[4])))
		elif gene_name == gene_2:
			exons_2.append(
				('chr' + tokens[0], tokens[6], int(tokens[3]), int(tokens[4])))
		
	gtf_file.close()

	for k, ex in enumerate(exons_1):
		seq = chromosomes[ex[0]][ex[2]-1:ex[3]]
		if ex[1] == '-': seq = revcomplement(seq)
		seq = seq[-50:]
		exons_1[k] = ex + (seq,)
	for k, ex in enumerate(exons_2):
		seq = chromosomes[ex[0]][ex[2]-1:ex[3]]
		if ex[1] == '-': seq = revcomplement(seq)
		seq = seq[:50]
		exons_1[k] = ex + (seq,)

	junctions_path = out_dir + '/junctions.fa'
	junctions_file = open(junctions_path, 'w')
	
	for ex1, ex2 in itertools.product(exons_1, exons_2):
		junctions_file.write('>%s[%d..%d]-%s[%d..%d]\n%s%s\n' %
			(gene_1, ex1[2], ex1[3], gene_2, ex2[2], ex2[3], ex1[4], ex2[4]))
	
	junctions_file.close()
	
	
	
	sample_name = os.path.basename(read_paths[0])
	sample_name = re.sub('(_[12])?\.(fastq|fasta|fq|fa).*', '', sample_name,
		flags=re.I)
	info('Starting fusion analysis on %s...' % sample_name)
	
	info('Discarding reads that align against the transcriptome...')
	txome_unaligned_path = out_dir + '/txome_unaligned.fq.gz'
	shell('bowtie2 -p8 --score-min L,0,-0.2 '
		'-x ~/tools/bowtie2-indexes/homo_sapiens/ensembl_68 --un-gz %s -U %s '
		'> /dev/null' % (txome_unaligned_path, ' '.join(read_paths)))
	
	chromosomes = read_fasta(
		'/data/csb/tools/bowtie2-indexes/homo_sapiens/hg19.fa')





	

#####################
# PHUSION VISUALIZE #
#####################

def visualize_fusions(fusion_report_paths):
	
	for s, report_path in enumerate(fusion_report_paths):
		report = open(report_path)
		for line in report:
			if not line.startswith('chr'): continue
			
			tokens = line[:-1].split('\t')
			print('%s:%d::\t%s:%d::' % 
				(tokens[0].replace('chr', ''), int(tokens[2]),
				tokens[5].replace('chr', ''), int(tokens[7])))
			
		report.close()



	
	

#######################
# COMMAND LINE PARSER #
#######################

if __name__ == '__main__':
	args = docopt.docopt(__doc__)
	if args['detect']:
		detect_fusions(args['<reads>'], out_dir=args['--output-dir'])
	elif args['blacklist']:
		create_fusion_blacklist(args['<fusion_report>'])
	elif args['filter']:
		filter_fusions(args['<fusion_report>'][0],
			region_blacklist_path=args['--blacklist'],
			min_flank_len=int(args['--min-flank-len']))
	elif args['rank']:
		rank_fusions(args['<fusion_report>'])
	elif args['validate']:
		validate_fusion(args['<fusion>'], args['<reads>'])
	elif args['visualize']:
		visualize_fusions(args['<fusion_report>'])
	
